import seaborn as sns
import matplotlib.pyplot as plt
import pandas as pd

lva = "close_var_l_ma_fib_m1"
fva = "close_var_f_ma_fib_m1"
va = "close_var_ma_fib_m1"
turn = 'turnover_ma_fib_m1'
cma = 'close_ma_fib_m1'
sns.set_style('ticks')
sns.set_context('paper')


class Strategy:
    pass


class StrategyVarMaFib(Strategy):
    def __init__(self):
        self.env = {}
        pass

    def policy(self, stock, row):
        score = self.update_env(stock, row)
        if row[lva + "upd"] > 0.07 and row[va + "upd"] > 0 \
                and row[fva + "upd"] > 0 and row[turn + "upd"] > 0 and row[cma + "upd"] > 0 and row["quote_rate"] < 0.05\
                and row["close_var_f_ma5"] > 0.15:
            # and row["close_var_f_ma20"] > 0.057
            # and row["close_var_f_ma60"] > 0.1
            return score, "buy"
        if row[lva + "upd"] < -0.06:
            return score, "sell"
        return score, ""

    def update_env(self, stock, row):
        if stock not in self.env:
            # 存3月， 价格浮动 与 lva的积分， 作为最终的
            self.env[stock] = {
                "score": 0
            }
        # 60日
        score = self.env[stock]["score"] + row[lva + "upd"] - 10 * row[cma + "upd"]
        self.env[stock]["score"] = max(0, round(score, 2))

        return self.env[stock]["score"]
class FinanceManage:

    def __init__(self, strategy: Strategy):
        self.init_cash = 1000000
        self.pool = {}
        self.loan_map = {}
        self.strategy = strategy
        self.return_map = {}


    def market_day(self, stock, row):
        if stock not in self.pool:
            self.pool[stock] = STrade(stock, self)
        score, action = self.strategy.policy(stock, row)
        self.pool[stock].action(row["close"], score, action, row["date"])

    def apply_cash(self):
        money = min(10000, self.init_cash)
        self.init_cash = self.init_cash - money
        return money

    def send_cash(self, stock, money, cost, rate):
        self.init_cash = self.init_cash + money
        if rate not in self.return_map:
            self.return_map[rate] = 1
        else:
            self.return_map[rate] = self.return_map[rate] + 1
        if stock not in self.loan_map:
            self.loan_map[stock] = 0
        self.loan_map[stock] = self.loan_map[stock] + money

    def loan(self, stock, money):
        if stock not in self.loan_map:
            self.loan_map[stock] = 0
        self.loan_map[stock] = self.loan_map[stock] + money

class STrade:
    stock = ""
    cost = 0
    cut_line = 0
    position = 0
    each_add_day = 2  # 多少天后可以加仓
    position_days = 0
    history = []
    finance_manage = None
    profit = 0

    def __init__(self, stock: str, finance_manage: FinanceManage):
        self.stock = stock
        self.finance_manage = finance_manage
        self.scat = {
            "day": [],
            "price": [],
            "action": [],
            "score": []
        }

    def action(self, price, score, action, day):
        if self.position > 0:
            self.cut_line = max(self.cut_line, price)
        if action == "buy":
            self.buy(price, score, day)
        elif action == "sell":
            self.sell(price, score, day)
        else:
            self.check_day(price, day)
        self.scat["score"].append(score)

        # 买入
    def buy(self, price, score, day):
        # 申请资金
        if self.cost > 30000:
            return
        if self.position == 0:
            self.do_buy(price, score, day)
        elif self.position_days % self.each_add_day == 0:
            self.do_buy(price, score, day)

    # 卖出
    def sell(self, price, score, day, act="sell"):
        if self.position > 0:
            money = self.position * price
            self.profit = money - self.cost
            rate = round(self.profit/self.cost, 2)
            self.finance_manage.send_cash(self.stock, money, self.cost, rate)
            print(self.stock, act, day, price, self.cost, rate)
            self.clear()
            self.scat["day"].append(day)
            self.scat["price"].append(price)
            self.scat["action"].append(act)
    # 观望
    def check_day(self, price, day):
        # 判断是否需要止损
        if self.position > 0:
            if self.cut_line * 0.75 > price:
                self.sell(price, 0, day, act="stop")
            else:
                self.position_days = self.position_days + 1
                # print(self.stock, "check", day, price, self.cut_line)

    def do_buy(self, price, score, day):
        money = self.finance_manage.apply_cash()
        self.position = self.position + (money / price)
        self.cost = self.cost + money
        self.cut_line = max(self.cut_line, price)
        self.position_days = self.position_days + 1
        self.scat["day"].append(day)
        self.scat["price"].append(price)
        self.scat["action"].append("buy")
        print(self.stock, "buy", day, price)

    def clear(self):
        self.position_days = 0
        self.cost = 0
        self.position = 0
        self.cut_line = 0

    # 交易画图
    def draw(self, mdata):
        f = plt.figure()
        f.add_subplot(3, 1, 1)
        plt.title(self.stock)
        sns.scatterplot(pd.Series(self.scat["day"]), pd.Series(self.scat["price"]),
                        hue=pd.Series(self.scat["action"]),
                        palette="Set2")
        sns.lineplot(data=mdata, x='date', y='close')
        f.add_subplot(3, 1, 2)
        plt.title("var_ma")
        sns.lineplot(data=[mdata[va], mdata[fva], mdata[lva], mdata['turnover_ma_fib_m1']])
        f.add_subplot(3, 1, 3)
        plt.title("var_upd")
        # sns.lineplot(data=[mdata["close_var_ma5_upd"], mdata["close_var_ma10_upd"]])
        dataScore = pd.DataFrame({"score": pd.Series(self.scat["score"], index=mdata["date"])})
        # sns.lineplot(data=dataScore)
        sns.lineplot(data=[mdata["close_var_l_ma60"], mdata['close_var_l_ma20'],
                           mdata["close_var_f_ma60"], mdata['close_var_f_ma20'],
                           ])
        # mdata["turnover_r60"], mdata['turnover_r20'],
        # plt.savefig("D:\\code\\rlearn\\picture1\\" + symbol + ".png")
        plt.show()
        xxx = 1
